# -*- coding:utf-8 -*-
"""
该脚本为灯光组关联脚本
"""
import sys
from PySide2.QtGui import *
from PySide2.QtCore import *
from PySide2.QtWidgets import *
try:

    import maya.cmds as cmds
    import maya.mel as mel
    from maya import OpenMayaUI as omui
    from shiboken2 import wrapInstance
except Exception as import_err:
    print('test')

def maya_main_window():
    """
     Return the Maya main window widget as a Python object
    :return:
    """
    main_window_ptr = omui.MQtUtil.mainWindow()
    return wrapInstance(long(main_window_ptr), QWidget)

class NameWindow(QWidget):
    name_Signal = Signal(str)
    def __init__(self,label_name):
        super(NameWindow, self).__init__()

        self.main_layout = QHBoxLayout()
        self.name_label = QLabel(label_name)
        self.name_edit = QLineEdit()
        self.btn = QPushButton(u'确认')
        self.btn.clicked.connect(self._ok)
        self.main_layout.addWidget(self.name_label)
        self.main_layout.addWidget(self.name_edit)
        self.main_layout.addWidget(self.btn)
        self.setLayout(self.main_layout)

    def _ok(self):
        text = self.name_edit.text()
        if text:
            self.name_Signal.emit(text)
            self.close()


class LightGroupList(QListWidget):

    clear_Signal = Signal(int)
    show_lights_Signal = Signal(list)
    def __init__(self):
        super(LightGroupList, self).__init__()
        self.setContextMenuPolicy(Qt.CustomContextMenu)  # 开启自定义右击菜单
        self.customContextMenuRequested.connect(self.right_menu)
        self.itemClicked.connect(self._show_lights)
        self.lg_data = {}  # {'lg01':[],'lg02':[],...,}

    def initLG(self,data):
        self.clear()
        self.lg_data = data
        self.addItems(data.keys())

    def _show_lights(self,item):
        text = item.text()
        lights_list = self.lg_data.get(text,[])
        self.show_lights_Signal.emit(lights_list)


    def _check_exist(self,name):
        result = False
        for i in range(self.count()):
            item = self.item(i)
            text = item.text()
            if text == name:
                result = True
        return result

    def addLG(self,name):
        if not self._check_exist(name):
            item = QListWidgetItem()
            item.setText(name)
            self.addItem(item)
            self.lg_data[name] = []
        else:
            QMessageBox.information(None,u'错误',u'当前灯光组已存在')

    def addLights(self,light_list):
        """
        添加灯光到当中选中的灯光组中
        :param light_list: 
        :return: 
        """
        item = self.currentItem()
        if not item:
            return

        for v in self.lg_data.values():  # 先清掉已有灯光
            for l in light_list:
                light_name = l.split('___')[-1]
                for name in v:
                    if light_name in name:
                        v.remove(name)

        current_lights = self.lg_data.get(item.text(),[]) # 再添加灯光
        for i in light_list:
            if i not in current_lights:
                current_lights.append(i)
                
        self.show_lights_Signal.emit(current_lights)  # 刷新灯光列表
                
    def removeLights(self,light_list):
        current_item = self.currentItem()
        current_lights = self.lg_data.get(current_item.text())
        for i in light_list:
            if i in current_lights:
                current_lights.remove(i)


    def right_menu(self):
        rMenu = QMenu()
        createAction = QAction(u'创建灯光组',self)
        createAction.triggered.connect(self._create_lg)
        rMenu.addAction(createAction)
        delAction = QAction(u'删除灯光组', self)
        delAction.triggered.connect(self._delete_lg)
        rMenu.addAction(delAction)
        clearAction = QAction(u'清空灯光组', self)
        clearAction.triggered.connect(self._clear_lg)
        rMenu.addAction(clearAction)
        rMenu.exec_(QCursor.pos())



    def _create_lg(self):
        self.name_window = NameWindow(u'灯光组')
        self.name_window.name_Signal.connect(self.addLG)
        self.name_window.show()

    def _delete_lg(self):
        item = self.currentItem()
        text = item.text()
        self.lg_data.pop(text)
        self.takeItem(self.currentRow())

    def _clear_lg(self):
        self.lg_data = {}
        self.clear()
        self.clear_Signal.emit(1)



class LightList(QListWidget):

    light_Signal = Signal(list)
    def __init__(self):
        super(LightList, self).__init__()
        self.setContextMenuPolicy(Qt.CustomContextMenu)  # 开启自定义右击菜单
        self.customContextMenuRequested.connect(self.right_menu)
        self.setSelectionMode(QAbstractItemView.ExtendedSelection)

    def right_menu(self):
        rMenu = QMenu()

        delAction = QAction(u'移除灯光', self)
        delAction.triggered.connect(self._delete_lg)
        rMenu.addAction(delAction)
        rMenu.exec_(QCursor.pos())

    def _delete_lg(self):
        items = self.selectedItems()
        light_list = []
        for i in items:
            text = i.text()
            self.takeItem(self.indexFromItem(i).row())
            light_list.append(text)
        self.light_Signal.emit(light_list)


class LightGroupWindow(QWidget):

    def __init__(self,parent=None):
        super(LightGroupWindow, self).__init__(parent)

        self.big_font = QFont('',14,65)
        self.mid_font = QFont('',11,65)
        # self.setParent(parent)
        self.resize(450,600)
        self.setWindowFlags(Qt.WindowStaysOnTopHint|Qt.Window)
        self.setWindowTitle(u'灯光组管理')
        self.main_layout = QVBoxLayout()
        self.tool_name = QLabel(u'灯光组管理器')
        self.tool_name.setFont(self.big_font)
        self.tool_name.setAlignment(Qt.AlignCenter)
        self.tool_name.setContentsMargins(0,10,0,20)

        self.center_layout = QHBoxLayout()
        self.center_layout.setSpacing(15)

        self.left_layout = QVBoxLayout()
        self.lg_label = QLabel(u'灯光组')
        self.lg_label.setContentsMargins(0,0,0,10)
        self.lg_label.setFont(self.mid_font)
        self.lg_label.setAlignment(Qt.AlignCenter)
        self.lg_layout = QHBoxLayout()
        self.lg_edit = QLineEdit()
        self.lg_edit.setText('LG1')
        self.lg_create_btn = QPushButton(u'创建灯光组')
        self.lg_create_btn.setShortcut('Return')
        self.lg_create_btn.clicked.connect(self._create_lg)
        self.lg_layout.addWidget(self.lg_edit)
        self.lg_layout.addWidget(self.lg_create_btn)
        self.lg_list = LightGroupList()
        self.lg_list.clear_Signal.connect(self._clear_lg)
        self.lg_list.show_lights_Signal.connect(self._show_lights)
        self.left_layout.addWidget(self.lg_label)
        self.left_layout.addLayout(self.lg_layout)
        self.left_layout.addWidget(self.lg_list)


        self.right_layout = QVBoxLayout()
        self.light_label = QLabel(u'灯光')
        self.light_label.setContentsMargins(0,0,0,10)
        self.light_label.setFont(self.mid_font)
        self.light_label.setAlignment(Qt.AlignCenter)
        self.light_add_btn = QPushButton(u'添加灯光')
        self.light_add_btn.clicked.connect(self._addToLG)
        self.light_list = LightList()
        self.light_list.light_Signal.connect(self._remove_lights)
        self.light_list.itemClicked.connect(self._select_lights)
        self.right_layout.addWidget(self.light_label)
        self.right_layout.addWidget(self.light_add_btn)
        self.right_layout.addWidget(self.light_list)

        self.center_layout.addLayout(self.left_layout)
        self.center_layout.addLayout(self.right_layout)

        self.tip_label = QLabel(u'注意：要进灯光组的原命名不能有___，还有不要在物体shape上手动修改灯光组')

        self.main_layout.addWidget(self.tool_name)
        self.main_layout.addLayout(self.center_layout)
        self.main_layout.addWidget(self.tip_label)
        self.setLayout(self.main_layout)

        self.init_LG()


    def init_LG(self):
        try:
            lg_data = cmds.getAttr('top.light_group')
            if lg_data:
                self.lg_list.initLG(eval(lg_data))
        except :
            print('test')

    def save_LG(self):
        try:
            cmds.addAttr('top', ln='light_group', dt='string')
        except:
            pass
        lg_data = self.lg_list.lg_data
        if lg_data:
            cmds.setAttr('top.light_group',str(lg_data),type='string')

    def increaseLgEdit(self):
        text = self.lg_edit.text() 
        if text:
            pre, num = text[:-1], text[-1:]
            try:
                new_num = int(num) + 1
                new_name = pre + str(new_num)
            except:
                new_name = text + '1'

            self.lg_edit.setText(new_name)

    def _create_lg(self):
        lg_name = self.lg_edit.text()
        if lg_name:
            self.lg_list.addLG(lg_name)
            self.increaseLgEdit()
        self.save_LG()

    def _clear_lg(self):
        self.light_list.clear()
        self.save_LG()

    def _show_lights(self,lights_list):
        self.light_list.clear()
        self.light_list.addItems(lights_list)
        try:
            cmds.select(lights_list)
        except :
            pass

    def _addToLG(self):
        sn = cmds.ls(sl=1)
        shapes = cmds.listRelatives(sn, shapes=1)
        lg_item = self.lg_list.currentItem()
        new_shapes = []
        if lg_item:
            lg_name = lg_item.text()
            for i in range(len(shapes)):
                try:
                    cmds.setAttr(shapes[i]+'.aiAov',lg_name,type='string')
                    # print(1111)
                    name = sn[i].split('___')[-1]
                    # print(22222,name)
                    new_name = lg_name+'___'+ name
                    # print(3333,new_name)
                    cmds.rename(sn[i],new_name)
                    # print(4444)
                    new_shape = cmds.listRelatives(new_name,shapes=True)[0]
                    # print(5555)
                    new_shapes.append(new_shape)
                except Exception as _addToLG_ERR:
                    print('_addToLG_ERR:',_addToLG_ERR)
            self.lg_list.addLights(new_shapes)
            self.save_LG()

    def _remove_lights(self,light_list):
        if light_list:
            self.lg_list.remove_lights(light_list)
        self.save_LG()

    def _select_lights(self):
        lights_items = self.light_list.selectedItems()
        if lights_items:
            lights_list = []
            for i in lights_items:
                light_shape = i.text()
                lights_list.append(light_shape)
            cmds.select(lights_list)

    def closeEvent(self, *args, **kwargs):
        self.save_LG()

def _run():
    try:
        dialog.show()
        dialog.setHidden(False)
    except :

        dialog = LightGroupWindow(parent=maya_main_window())
        dialog.show()

def _doc():
    doc = {'name':u'灯光组管理','doc':u'管理灯光组'}
    return doc

if __name__ == '__main__':
    app = QApplication(sys.argv)
    l = LightGroupWindow(None)
    l.show()
    sys.exit(app.exec_())
        






